class Trainer:
    def __init__(self, model, criterion, optimizer):
        self.model = model
        self.criterion = criterion
        self.optimizer = optimizer

    def train(self, dataloader, epochs):
        self.model.train()
        for epoch in range(epochs):
            total_loss = 0.0
            for src, tgt, y in dataloader:
                self.optimizer.zero_grad()
                y_hat = self.model(src, tgt)
                loss = self.criterion(y, y_hat)
                loss.backward()
                self.optimizer.step()
                total_loss += loss.item()
            print(f"Epoth {epoch}, Training loss: {total_loss / len(dataloader)}")
